#!/usr/bin/env python
# _*_ coding: utf-8 _*_
#
#
# Copyright (C) 2020 Wei Keting<weikting@gmail.com>. All rights reserved.
# @Time : 2020-09-10 14:01 
# @File : visualize.py.py
# @Description :
#
#


import os

import matplotlib.pyplot as plt
import numpy as np
import time

# noinspection PyUnresolvedReferences
from mpl_toolkits.mplot3d.axes3d import Axes3D

COLORS = ['#d216a5', '#606100', '#438e3d', '#0226b9', '#282473', '#007767', '#94f463', '#a0deff', '#bb0c1f', '#0895f1']


def visualize(feat, labels, title, filename=None, output_dir=None, s=5, **kwargs):
    # plt.ion()
    global COLORS
    c = COLORS.copy()
    c_labels = len(np.unique(labels))
    n_c = len(c)
    if c_labels > n_c:
        for i in range(n_c, c_labels):
            c.append(COLORS[i % 10])
    labels = np.squeeze(labels)
    plt.figure(figsize=(12, 8))
    plt.clf()
    plt.grid(b=True, axis='x', alpha=0.5)  # 只显示x轴网格线
    plt.grid(b=True, axis='y', alpha=0.5)  # 只显示y轴网格线
    for i in range(len(c)):
        plt.scatter(feat[labels == i, 0], feat[labels == i, 1], c=c[i], s=2, alpha=0.8)

    plt.title(title)
    plt.xlabel("X")
    plt.ylabel("Y")
    plt.legend(['{}'.format(x) for x in range(c_labels)], loc='upper right', markerscale=s)

    extras = kwargs.get('extras', None)
    if extras is not None:
        for w in extras:
            nw = w / np.linalg.norm(w)
            if nw[0] >= 0:
                x = np.arange(0, max(w[0], nw[0]), 0.1)
            else:
                x = np.arange(0, min(w[0], nw[0]), -0.1)
            plt.plot(x, x * w[1] / w[0], c='grey', linewidth=1, alpha=0.85)
            plt.scatter(w[0], w[1], s=24, c='black')
            plt.scatter(w[0] * 2, w[1] * 2, s=24, c='white')
            plt.scatter(w[0] * 2, w[1] * 2, s=48, c='red')

    if output_dir is None:
        output_dir = os.path.expanduser("~/tmp")
    pngdir = os.path.join(output_dir, 'visualizations')
    os.makedirs(pngdir, exist_ok=True)

    if filename is None or len(filename) == 0:
        filename = "{}-{}.png".format(title, time.strftime("%Y%m%d-%H%M%S"))
    out_file = os.path.join(pngdir, filename)
    plt.savefig(out_file, dpi=160)
    print("visualize_2d: ", out_file)

    # 显示图例
    if kwargs.get('show', False):
        plt.show()
    plt.close()


def visualize_model(model, datas, title, max_count=0, **kwargs):
    fts = None
    labels = None
    t = time.time()
    for d, l in datas:
        f = model.predict(d)
        if fts is None:
            fts = f
            labels = l
        else:
            fts = np.concatenate((fts, f))
            labels = np.concatenate((labels, l))
        if 0 < max_count <= len(fts):
            break
    t = time.time() - t
    print("fts: ", fts.shape, "labels: ", labels.shape)
    if fts.shape[-1] >= 3:
        fts = fts / np.clip(np.linalg.norm(fts, axis=-1, keepdims=True), 1e-7, 1.1)
    print("norm fts: ", fts.shape, "labels: ", labels.shape)
    print("predict time: ", t, "seconds", t * 1000. / len(fts), " ms per sample")
    visualize(fts, labels, title, **kwargs)
